# author：曹一鸣
import winreg
import json


class Vulnerabilty():
    # 1. 查找是否存在影子账户
    @staticmethod
    def check_shadow_account():
        root = winreg.ConnectRegistry(None, winreg.HKEY_LOCAL_MACHINE)
        try:
            # 查看注册表中的账户
            reg_path = r"SAM\\SAM\\Domains\\Account\\Users"
            key = winreg.OpenKey(root, reg_path)
            results = {}
            byteList = []
            try:
                count = 0
                while(1):
                    try:
                        # 查看全部账户并保存
                        f = winreg.EnumKey(key, count)
                        results[count] = f
                        count += 1
                    except EnvironmentError:
                        break
            finally:
                winreg.CloseKey(key)

            for i in range(len(results)-1):
                reg_path = r"SAM\\SAM\\Domains\\Account\\Users"+"\\"+results[i]
                key = winreg.OpenKey(root, reg_path)
                try:
                    count = 0
                    while(1):
                        try:
                            regKey, regValue, regType = winreg.EnumValue(key, count)
                            # 查看账户的键值并保存
                            if regKey == 'F':
                                byteList.append(regValue)  
                                break
                            count += 1
                        except EnvironmentError:
                            break
                finally:
                    winreg.CloseKey(key)
        finally:
            winreg.CloseKey(root)

        new_byteList = list(dict.fromkeys(byteList))

        # 比较是否存在重复键值，存在则说明存在影子账户
        if len(new_byteList) == len(byteList):
            info = {
                "shadowAccount":"not found"
            }
            
        else:
            info = {
                "shadowAccount":"exist"
            }
        return json.dumps(info, ensure_ascii=False)

    # 2. 查找映像劫持
    @staticmethod
    def check_image_hijack():
        root = winreg.ConnectRegistry(None, winreg.HKEY_LOCAL_MACHINE)
        try:
            reg_path = r"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Image File Execution Options"
            key = winreg.OpenKey(root, reg_path)
            results = {}
            imageHijackList = []
            try:
                count = 0

                while(1):
                    try:
                        f = winreg.EnumKey(key, count)
                        results[count] = f
                        count += 1
                    except EnvironmentError:
                        break
            # finally是无论是否有异常，最后都要做的一些事情
            finally:
                winreg.CloseKey(key)

            for i in range(len(results)):
                reg_path = r"SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Image File Execution Options"+"\\"+results[i]
                try:
                    key = winreg.OpenKey(root, reg_path)
                except EnvironmentError:
                    continue
                try:
                    count = 0
                    while(1):
                        try:
                            # 查看键值里是否有Debugger
                            regKey,regValue,regType = winreg.EnumValue(key, count)
                            if regKey == 'Debugger':
                                imageHijackList.append(results[i])  
                            count += 1
                        except EnvironmentError:
                            break
                # finally是无论是否有异常，最后都要做的一些事情
                finally:
                    winreg.CloseKey(key)
        # finally是无论是否有异常，最后都要做的一些事情
        finally:
            winreg.CloseKey(root)

        keys = [str(x+1) for x in range(len(imageHijackList))]
        imageHijackList_json = dict(zip(keys, imageHijackList))
        return json.dumps(imageHijackList_json, ensure_ascii=False)


def main():
    test = Vulnerabilty()
    print(test.check_shadow_account())
    print(test.check_image_hijack())

if __name__ == '__main__':
    main()
